import os
import cv2
import numpy as np
import open3d as o3d
import argparse
from utils import read_d
import matplotlib.pyplot as plt
from functools import partial

def rotate_view(vis, path_json=""):
    ctr = vis.get_view_control()
    if path_json:
        parameters = o3d.io.read_pinhole_camera_parameters(path_json)
        ctr.convert_from_pinhole_camera_parameters(parameters)
    return False


def depth2pcd(args):
    if args.calib_file:
        cv_file = cv2.FileStorage(args.calib_file, cv2.FILE_STORAGE_READ)    
        f_disp = cv_file.getNode("proj_matL").mat()[0,0]
        cx = cv_file.getNode("proj_matL").mat()[0,2]
        cy = cv_file.getNode("proj_matL").mat()[1,2]
        b_disp = float(cv_file.getNode("baselineLR").real()) / 1000.

        cv_file.release()
    else:
        f_disp = args.f
        cx = args.cx
        cy = args.cy
        b_disp = args.baseline

    rgb = cv2.imread(args.rgb_path).astype(np.uint8)
    rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
    
    depth = read_d(args.depth_path, invert=args.is_disparity, f=f_disp, b=b_disp, scale_factor=args.scale_factor)
    
    if args.path_conf:
        conf = np.load(args.path_conf)
    else:
        conf = np.ones_like(depth)


    
    # Getting shapes and calculating scale factor
    factor = rgb.shape[0] / depth.shape[0] 
    # Resize depth to match RGB shape
    depth = cv2.resize(depth, (rgb.shape[1], rgb.shape[0]), interpolation=cv2.INTER_NEAREST) / factor

    # Creating mask
    mask = 1 - (np.isinf(depth)).astype(np.float32)
    conf = cv2.resize(conf, (rgb.shape[1], rgb.shape[0]))
    th = np.percentile(conf, args.th)
    mask[conf>th] = 0
    breakpoint()
    #!==============================
    # depth = depth * mask    
    depth[(1 - mask).astype(np.bool_)] =0
    #!==============================


    # Cropping
    crop_left = args.crop_left
    crop_right = args.crop_right
    crop_top = args.crop_top
    crop_bottom = args.crop_bottom
    rgb = rgb[crop_top:rgb.shape[0]-crop_bottom,crop_left:rgb.shape[1]-crop_right,:].copy()
    depth = depth[crop_top:depth.shape[0]-crop_bottom,crop_left:depth.shape[1]-crop_right].copy()
    cx = cx - crop_left
    cy = cy - crop_top
    h, w = rgb.shape[:2]

    if not args.not_visualize_images:
        plt.subplot(1,2,1)
        plt.imshow(rgb)
        plt.title("Image-Left")
        plt.subplot(1,2,2)
        plt.imshow(depth)
        plt.title("Depth-Left (meters)")
        plt.show()

    depth = o3d.geometry.Image(depth)
    rgb = o3d.geometry.Image(rgb)
    rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(rgb, 
                                                            depth, 
                                                            depth_scale=1., 
                                                            depth_trunc=args.max_depth, 
                                                            convert_rgb_to_intensity=False)
                                                    
    intrinsic = o3d.camera.PinholeCameraIntrinsic(w, h, f_disp, f_disp, cx, cy)
    point_cloud = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, intrinsic)

    if not args.not_visualize_pcl:
        rotate_view_part = partial(rotate_view, path_json=args.path_json)
        o3d.visualization.draw_geometries_with_animation_callback([point_cloud], rotate_view_part)

    if args.save_ply:
        os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
        o3d.io.write_point_cloud(args.output_path, point_cloud)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Argument parser')
    parser.add_argument('--rgb_path', type=str, help='left img path')
    
    parser.add_argument('--depth_path', type=str, help='left img path')
    parser.add_argument('--is_disparity', action='store_true', help='disparity path')
    parser.add_argument('--max_depth', type=float, default=20, help='clip depth to this value')
    parser.add_argument('--scale_factor', type=int, default=32, help='scale factor used to decode 16 bit depth png images')

    parser.add_argument('--path_conf', dest='path_conf', type=str, default='', help='path conf')
    parser.add_argument('--th', type=float, help='conf threshold', default=100)
    
    parser.add_argument('--crop_bottom', type=int, help='crop bottom', default=0)
    parser.add_argument('--crop_top', type=int, help='crop top', default=0)
    parser.add_argument('--crop_left', type=int, help='crop left', default=0)
    parser.add_argument('--crop_right', type=int, help='crop right', default=0)

    parser.add_argument('--calib_file', type=str, default='', help='calib file used to convert disparity to depth. If none uses f, cx, cy and baseline args')
    parser.add_argument('--f', type=float, help='f', default=4.6475097656250000e+03)
    parser.add_argument('--cx', type=float, help='cx', default=2.0263100128173828e+03)
    parser.add_argument('--cy', type=float, help='cy', default=1.5531697540283203e+03)
    parser.add_argument('--baseline', type=float, help='cy', default=0.082142385598336574)

    parser.add_argument('--save_ply',  help='save_ply', action='store_true')
    parser.add_argument('--output_path', type=str, default='pcd.ply', help='output path')

    parser.add_argument('--path_json', type=str, default='', help='path_json of viewpoint')

    parser.add_argument('--not_visualize_images',  help='not visualize 2D images', action='store_true')
    parser.add_argument('--not_visualize_pcl',  help='not visualize pcl', action='store_true')

    args = parser.parse_args()
    depth2pcd(args)